package com.justgame.project.cloud.gateway.config;

import com.alibaba.fastjson.JSONObject;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

/**
 * @author ZhouWeiBo
 * @create 2022/4/22
 */
public class GatewayResponseDecorator extends ServerHttpResponseDecorator {

    public GatewayResponseDecorator(ServerHttpResponse delegate) {
        super(delegate);
    }

    @Override
    @SuppressWarnings(value = "unchecked")
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        if(body instanceof Flux) {
            Flux<DataBuffer> fluxBody = (Flux<DataBuffer>) body;
            return super.writeWith(fluxBody.buffer().map(dataBuffers -> {
                DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                DataBuffer join = dataBufferFactory.join(dataBuffers);
                byte[] content = new byte[join.readableByteCount()];
                join.read(content);
                DataBufferUtils.release(join);// 释放掉内存
                String bodyStr = new String(content, StandardCharsets.UTF_8);
                //修改响应体
                bodyStr = modifyBody(bodyStr);
                getDelegate().getHeaders().setContentLength(bodyStr.getBytes().length);
                return bufferFactory().wrap(bodyStr.getBytes());
            }));
        }
        return super.writeWith(body);
    }

    //重写这个函数即可
    private String modifyBody(String jsonStr){
        //JSONObject json = JSONObject.parseObject(jsonStr);
        //TODO...修改响应体
        return jsonStr;
    }
}
